import torch
import torch.nn as nn

from transcoder.models.block import TRANS_BLOCK

class BeamformModel(nn.Module):
    def __init__(self, K, Nc, Nt, Q, **kwargs):
        super().__init__()
        self.K = K
        self.Nc = Nc
        self.Nt = Nt
        self.Q = Q
        self.mmfuse = TRANS_BLOCK(2*self.K+64, 2*self.K, 128, 4)  #数据驱动预编码
        self.DDDL = TRANS_BLOCK(2*self.K+64, 2*self.Nt, 128, 4)  #数据驱动预编码

        self.beta1 = nn.Parameter(torch.zeros(1))
        self.beta2 = nn.Parameter(torch.ones(1))
    
    def forward(self, qsymbol, HH):
        #################### 图像和通信语义融合 #######################
        qsymbol = qsymbol.reshape(-1, self.Nc, 2, self.Q) # [B*K, Nc, 2, Q]
        
        qsymbol = qsymbol.reshape(-1, self.K, self.Nc, 2, self.Q) # [B, K, Nc, 2, Q]
        qsymbol = qsymbol.permute(0, 4, 2, 1, 3) # [B, Q, Nc, K, 2]
        
        B = qsymbol.shape[0]
        
        # 第Q个符号，在第Nc个载波上，发给K个用户
        qsymbol = qsymbol.reshape(-1, self.Nc, 2*self.K) # [B*Q, Nc, 2K]
        
        # qsymbol = qsymbol.reshape(-1,Nc,K*2*Q) # [4,64,K*2*Q]

        HH = HH.reshape(-1, 1, self.Nc, 64) # [B, 1, Nc, D]
        HH = HH.repeat(1, self.Q, 1, 1)     # [B, Q, Nc, D]
        HH = HH.reshape(-1, self.Nc, 64)    # [B*Q, Nc, D]
        
        z = torch.cat((qsymbol, HH), 2) # [B*Q, Nc, 2*K+D]

        z_fuse = self.mmfuse(z) # [B*Q, Nc, 2*K]

        z_ini = z_fuse + 0 # [B*Q, Nc, 2*K]
        # z_fuse = z_fuse.reshape(-1, self.Q, self.Nc, self.K, 2) # [B, Q, Nc, K, 2]
        # z_fuse = z_fuse.permute(0, 2, 3, 4, 1)   # [B, Nc, K, 2, Q]
        # z_model = z_fuse[:,:,:,0,:] + 1j * z_fuse[:,:,:,1,:] # [B, Nc, K, Q]

        ################### 数据驱动预编码 #######################
        # 利用CSI计算预编码矩阵，优化信号传输，即将3D语义特征转换成3D发射信号矩阵
        # 根据特征获得power，进行Beamforming
        signal = torch.cat((z_ini, HH), 2)     # [B*Q, 64, 2*K+64]
        signal = self.DDDL(signal) * self.beta1[0]  # [B*Q, 64, 2*Nt]
        signal = signal.reshape(B, self.Q, self.Nc, self.Nt, 2)    # [B, Q, Nc, Nt, 2]
        signal = signal.permute(0,2,3,4,1)          # [B, Nc, Nt, 2, Q]

        signal = signal[:,:,:,0,:] + 1j * signal[:,:,:,1,:] # [B, Nc, Nt, Q]
        signal = signal * self.beta2[0]
        
        return signal